from Model_Chooser import *
from Model import LinearModel



def train_all_model(k, cluster):
    for i in range(k):
        if len(cluster[i]) == 0:
            continue

        train_one_model(i, cluster[i])


def load_model(index):
    mdl = LinearModel(1, 1)
    mdl.load_state_dict(torch.load(r"D:\PySpace\Learned_Index\models\model{}.pkl".format(index)))
    return mdl


def get_data():
    with open("D:\\ep\\data.txt", "r") as f:
        lines = f.readlines()
    dataset = []
    for line in lines:
        k = line.split()[0]
        v = line.split()[1]
        dataset.append([int(k), int(v)])
    dataset = np.array(dataset)
    return dataset


if __name__ == "__main__":
    k = 10
    dataset = get_data()
    classifier = Classifier(k)
    #classifier.init_clusters(dataset)
    classifier.read_from_mem()
    #train_all_model(k, classifier.get_cluster())
    index = classifier.insert_key_value(1002, 1002)
    print(index)
    test_model = load_model(index)
    x = torch.tensor(1002, dtype=torch.float).view(1, -1)
    print(int(test_model(x)))

"""
10
52.5
157.5
261.0
363.5
464.5
564.0
663.0
761.0
857.5
952.5
106
104
103
102
100
99
99
97
96
94
"""